# ============================================================
# 该脚本用于自动识别衣服图片中的正反面，并将正反面配对拼接保存
# ============================================================

f'''
1. 解压zip得到images（本脚本已直接从_img_clo目录读取图片，无需解压）
2. 遍历images:
    - 如果是正面，fronts.append(img)
    - 如果是背面，backs.append(img)
    - 否则跳过
    - 如果fronts和backs都不为空：
        - 如何找到正面和背面配对？
        - 拼接[front, back]，判断是否为同一件衣服的正反面
        - 若是，则保存拼接图并跳出
逻辑更改：
for n in names:
    if not n endswith .jpg:
        continue
    n_noext, ext = os.path.splitext(n)
    n_split = n_noext.split('_')
    if n_split[-1] is digital:
        if n_split[-2] is digital:
            key = '_'.join( n_split[:-2] )
        else:
            key = '_'.join( n_split[:-1] )
        dict[key] .append osj( data_dir , n ) # 将文件名前缀相同的 路径放在一起

for k,v in dict:
    if v len < 2:continue
    front = []
    back = []
    for path in v:
        ans_front from path
        ans_back from path
        if check front:
            front append (path,pil)
        elif check back:
            fron append (path,pil)
        else:pass

        if front and back: # back 数量 一般比 front 少
            for b in back:
                for f in front:
                    f_path, f_pil = f
                    b_path, b_pil = b

                    f_name = osb(f_path)
                    b_name = osb(b_path)
                    
                    concat_img = concat(f_pil, b_pil)
                    save_name = f'f_name__b_name.jpg'
                    save_path = osj( save_dir, save_name )
                    if check ans is front_back:
                        if check ans is same color:
                            concat_img.save( save_path )

                            break
'''

import os
import shutil
import pdb
import zipfile
from PIL import Image
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

'''
改动说明：
- 直接从_img_clo目录获取图片
- 对图片名排序后，观察后10件衣服是否存在正反面配对
'''

from pathlib import Path
import cv2
import numpy as np
from utils.util_for_os import osj, ose, osb

# 获取数据和保存目录的lambda函数
get_data_dir = lambda category, date: f'/mnt/nas/datasets/diction/{category}{date}_img_clo'  # sifted clo
get_save_dir = lambda category, date: f'/mnt/nas/datasets/diction/{category}{date}_img_front_back_direct'  # img

# 只处理sweater类别和指定日期
categories = ['sweater']
dates = ["0916"]
# dates = ["0905", '0908', "0909"]

# 检查回答是否为否/是的辅助函数
check_no = lambda ans: 'no' in ans.lower() or '否' in ans or '不是' in ans
check_yes = lambda ans: 'yes' in ans.lower() or '是' in ans

import torch
from transformers import AutoConfig, AutoModelForCausalLM
from janus.models import VLChatProcessor

# 加载Janus多模态大模型和处理器
model_path = "/mnt/nas/zhangshu/Janus-Pro-7B"
config = AutoConfig.from_pretrained(model_path)
language_config = config.language_config
language_config._attn_implementation = 'eager'
vl_gpt = AutoModelForCausalLM.from_pretrained(
    model_path,
    language_config=language_config,
    trust_remote_code=True
)
if torch.cuda.is_available():
    vl_gpt = vl_gpt.to(torch.bfloat16).cuda()
else:
    vl_gpt = vl_gpt.to(torch.float16)

vl_chat_processor = VLChatProcessor.from_pretrained(model_path)
tokenizer = vl_chat_processor.tokenizer
cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu'

@torch.inference_mode()
def multimodal_understanding(image, question, seed=42, top_p=0.95, temperature=0.1):
    """
    使用Janus模型对图片和问题进行多模态理解，返回模型回答
    """
    # 清理CUDA缓存
    torch.cuda.empty_cache()

    # 设置随机种子，保证结果可复现
    torch.manual_seed(seed)
    np.random.seed(seed)
    torch.cuda.manual_seed(seed)

    # 构造对话输入
    conversation = [
        {
            "role": "<|User|>",
            "content": f"<image_placeholder>\n{question}",
            "images": [image],
        },
        {"role": "<|Assistant|>", "content": ""},
    ]

    pil_images = [Image.fromarray(image)]
    prepare_inputs = vl_chat_processor(
        conversations=conversation, images=pil_images, force_batchify=True
    ).to(cuda_device, dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float16)

    inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)

    outputs = vl_gpt.language_model.generate(
        inputs_embeds=inputs_embeds,
        attention_mask=prepare_inputs.attention_mask,
        pad_token_id=tokenizer.eos_token_id,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        max_new_tokens=512,
        do_sample=False if temperature == 0 else True,
        use_cache=True,
        temperature=temperature,
        top_p=top_p,
    )

    answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
    return answer

def get_answer(filepath, q, img_pil=None):
    """
    获取图片的模型回答
    - filepath: 图片路径
    - q: 问题
    - img_pil: 可选，已加载的PIL图片
    """
    if img_pil is not None:
        np_image = np.array(img_pil)  # shape: (H, W, 3)
        # RGB转BGR（OpenCV默认格式）
        cv2_image = cv2.cvtColor(np_image, cv2.COLOR_RGB2BGR)
    else:
        cv2_image = cv2.imread(filepath)
        if cv2_image is None:
            # 错误处理：图片路径无效或图片损坏
            print(f"Error: Cannot read image {filepath}")
            return ""
    return multimodal_understanding(cv2_image, q)

# 问题定义
Q8 = 'Is a front view from cloth'
Q9 = 'Is a back view from cloth'
def is_clo_front(ans_front, ans_back):
    """
    判断是否为衣服正面
    - ans_front: 回答Q8
    - ans_back: 回答Q9
    """
    if check_yes(ans_front) and check_no(ans_back):
        return True
    return False

def is_clo_back(ans_front, ans_back):
    """
    判断是否为衣服背面
    - ans_front: 回答Q8
    - ans_back: 回答Q9
    """
    if check_no(ans_front) and check_yes(ans_back):
        return True
    return False

Q10 = 'Is a front view and a back view from the same cloth'
Q11 = "What's the color of the left clo? What's the color of the right clo? Is the same color between the left and the right ?"
def check_Q10(ans):
    """
    判断拼接图像是否为同一件衣服的正反面
    """
    if check_yes(ans):
        return True
    return False

def check_Q11(ans):
    """
    判断左右衣服颜色是否相同
    """
    if 'same' in ans.lower() and 'the same color is not' not in ans.lower():
        return True
    return False

from tqdm import tqdm
from utils.util_flux import process_img_1024, horizontal_concat_images

# 1. 检查数据目录和names.txt文件是否存在
for category in categories:
    for date in dates:
        if date >= '0826' and category == 'leather':
            continue
        data_dir = get_data_dir(category, date)
        name_path = osj(data_dir, 'names.txt')
        assert ose(name_path), f"names.txt not found: {name_path}"

# 2. 主流程：遍历图片，识别正反面并配对保存
for category in categories:
    for date in dates:
        if date >= '0826' and category == 'leather':
            continue

        data_dir = get_data_dir(category, date)
        name_path = osj(data_dir, 'names.txt')
        assert ose(name_path), f"names.txt not found: {name_path}"
        with open(name_path, encoding='utf-8') as f:
            jpg_names = f.readlines()
        # 只保留.jpg结尾的文件名
        jpgnames = [name.strip() for name in jpg_names if name.strip().endswith('.jpg')]

        save_dir = get_save_dir(category, date)
        if ose(save_dir):
            shutil.rmtree(save_dir)
        os.makedirs(save_dir)

        # 注意：应对jpgnames排序，保证图片顺序一致
        jpgnames.sort()

        '''
        逻辑说明：
        - front/back为滑动窗口，最多保留3个
        - 遍历图片，判断正面/背面，分别加入front/back
        - front和back都不为空时，尝试两两配对
            - 拼接图片，问模型是否为同一件衣服的正反面
            - 若是，再判断左右颜色是否一致
            - 若都满足，则保存拼接图，并从front/back中移除已配对项
        '''
        front = []
        back = []
        for name in tqdm(jpgnames):
            path = osj( data_dir , name ) 

            # 问模型是否为正面/背面
            ans_front = get_answer(path, Q8)
            ans_back = get_answer(path, Q9)

            # 判断正面/背面，加入对应队列，队列最多3个
            if is_clo_front(ans_front, ans_back):
                if len(front) >= 3: front.pop(0)
                front.append((path, process_img_1024(path)))
            elif is_clo_back(ans_front, ans_back):
                if len(back) >= 3: back.pop(0)
                back.append((path, process_img_1024(path)))
            
            # front和back都不为空时，尝试配对
            if front and back:
                for f in reversed(front):
                    f_path, f_pil = f
                    for b in reversed(back):
                        b_path, b_pil = b
                        
                        front_name = osb(f_path)
                        back_name = osb(b_path)

                        # 拼接正反面图片
                        concat_img = horizontal_concat_images([f_pil, b_pil])
                        # 问模型：是否为同一件衣服的正反面
                        ans = get_answer('', Q10, img_pil=concat_img)
                        if check_Q10(ans):
                            # 再问模型：左右颜色是否一致
                            ans_color = get_answer('', Q11, img_pil=concat_img)
                            if check_Q11(ans_color):
                                # 保存拼接图像
                                concat_save_name = f'{front_name}__{back_name}.jpg'
                                concat_save_path = osj(save_dir, concat_save_name)
                                concat_img.save(concat_save_path)
                                concat_img.save('tmp_concat2.jpg')
                                print('concat img find !!!!!', flush=True)
                                print(concat_save_name, ans_color, flush=True)

                                # 找到正反面配对后，从队列中移除，跳出当前循环
                                front.remove(f)
                                back.remove(b)
                                break
